{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook explores using xarray's pointwise indexing, which allows you to extract cell values from a datacube for individual (x,y,x) tuples.  This is particularly useful for extracting values from a gridded data set (e.g. reanalysis) for a shiptrack, flightline or drifting station directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For starters I'm just going to set up a simple example using a 3-D array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4, x: 3, y: 4)>\n",
       "array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14, 15],\n",
       "        [16, 17, 18, 19],\n",
       "        [20, 21, 22, 23]],\n",
       "\n",
       "       [[24, 25, 26, 27],\n",
       "        [28, 29, 30, 31],\n",
       "        [32, 33, 34, 35]],\n",
       "\n",
       "       [[36, 37, 38, 39],\n",
       "        [40, 41, 42, 43],\n",
       "        [44, 45, 46, 47]]])\n",
       "Coordinates:\n",
       "  * t        (t) int64 0 1 2 3\n",
       "  * x        (x) int64 0 1 2\n",
       "  * y        (y) <U1 'a' 'b' 'c' 'd'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da = xr.DataArray(np.arange(4*12).reshape((4,3, 4)), dims=['t', 'x', 'y'], coords={'t': [0,1,2,3], 'x': [0, 1, 2], 'y': ['a', 'b', 'c', 'd']})\n",
    "da"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To utilize pointwise indexing, indices or labels have to be in DataArrays.  Here I'm setting up index arrays that I will use with isel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "it = xr.DataArray([0,1,2,3], dims='t')\n",
    "ix = xr.DataArray([0,1,2,2], dims='t')\n",
    "iy = xr.DataArray([0,1,3,2], dims='t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just use isel in the normal way.  This returns a 4 element DataArray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4)>\n",
       "array([ 0, 17, 35, 46])\n",
       "Coordinates:\n",
       "  * t        (t) datetime64[ns] 2018-08-15 2018-08-16 2018-08-17 2018-08-18\n",
       "    x        (t) int64 0 1 2 2\n",
       "    y        (t) int64 0 1 3 2"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da.isel(t=it, x=ix, y=iy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we just used lists or numpy arrays we get the full DataArray back because the index arrays span each dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4, x: 4, y: 4)>\n",
       "array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11],\n",
       "        [ 8,  9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14, 15],\n",
       "        [16, 17, 18, 19],\n",
       "        [20, 21, 22, 23],\n",
       "        [20, 21, 22, 23]],\n",
       "\n",
       "       [[24, 25, 26, 27],\n",
       "        [28, 29, 30, 31],\n",
       "        [32, 33, 34, 35],\n",
       "        [32, 33, 34, 35]],\n",
       "\n",
       "       [[36, 37, 38, 39],\n",
       "        [40, 41, 42, 43],\n",
       "        [44, 45, 46, 47],\n",
       "        [44, 45, 46, 47]]])\n",
       "Coordinates:\n",
       "  * t        (t) int64 0 1 2 3\n",
       "  * x        (x) int64 0 1 2 2\n",
       "  * y        (y) <U1 'a' 'b' 'c' 'd'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da.isel(t=[0,1,2,3], x=[0,1,2,2], y=[0,1,2,3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also access elements by label, for example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4)>\n",
       "array([ 0, 17, 34, 47])\n",
       "Coordinates:\n",
       "  * t        (t) int64 0 1 2 3\n",
       "    x        (t) int64 0 1 2 2\n",
       "    y        (t) <U1 'a' 'b' 'c' 'd'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dy = xr.DataArray(['a','b','c','d'], dims='t')\n",
    "da.sel(t=it, x=ix, y=dy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if we use timestamps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatetimeIndex(['2018-08-15', '2018-08-16', '2018-08-17', '2018-08-18'], dtype='datetime64[ns]', freq='D')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time = pd.date_range('2018-08-15', periods=4)\n",
    "time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4, x: 3, y: 4)>\n",
       "array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14, 15],\n",
       "        [16, 17, 18, 19],\n",
       "        [20, 21, 22, 23]],\n",
       "\n",
       "       [[24, 25, 26, 27],\n",
       "        [28, 29, 30, 31],\n",
       "        [32, 33, 34, 35]],\n",
       "\n",
       "       [[36, 37, 38, 39],\n",
       "        [40, 41, 42, 43],\n",
       "        [44, 45, 46, 47]]])\n",
       "Coordinates:\n",
       "  * t        (t) datetime64[ns] 2018-08-15 2018-08-16 2018-08-17 2018-08-18\n",
       "  * x        (x) int64 0 1 2\n",
       "  * y        (y) <U1 'a' 'b' 'c' 'd'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da = xr.DataArray(np.arange(4*12).reshape((4,3, 4)), dims=['t', 'x', 'y'], coords={'t': time, 'x': [0, 1, 2], 'y': ['a', 'b', 'c', 'd']})\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = xr.DataArray([dt.datetime(2018,8,16), \n",
    "                  dt.datetime(2018,8,16), \n",
    "                  dt.datetime(2018,8,18), \n",
    "                  dt.datetime(2018,8,18)], \n",
    "                 dims=['t'])\n",
    "ix = xr.DataArray([0,1,2,2], dims='t')\n",
    "iy = xr.DataArray([0,1,3,2], dims='t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4)>\n",
       "array([12, 17, 47, 46])\n",
       "Coordinates:\n",
       "  * t        (t) datetime64[ns] 2018-08-16 2018-08-16 2018-08-18 2018-08-18\n",
       "    x        (t) int64 0 1 2 2\n",
       "    y        (t) int64 0 1 3 2"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da.sel(t=t, x=ix, y=iy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Need to figure out how to use method='nearest' for times that are  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4)>\n",
       "array(['1979-08-16T00:00:00.000000000', '2018-08-16T00:00:00.000000000',\n",
       "       '2018-08-18T00:00:00.000000000', '2018-08-19T00:00:00.000000000'],\n",
       "      dtype='datetime64[ns]')\n",
       "Dimensions without coordinates: t"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = xr.DataArray([dt.datetime(1979,8,16), \n",
    "                  dt.datetime(2018,8,16), \n",
    "                  dt.datetime(2018,8,18), \n",
    "                  dt.datetime(2018,8,19)], \n",
    "                 dims=['t'])\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4, x: 3, y: 4)>\n",
       "array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14, 15],\n",
       "        [16, 17, 18, 19],\n",
       "        [20, 21, 22, 23]],\n",
       "\n",
       "       [[24, 25, 26, 27],\n",
       "        [28, 29, 30, 31],\n",
       "        [32, 33, 34, 35]],\n",
       "\n",
       "       [[36, 37, 38, 39],\n",
       "        [40, 41, 42, 43],\n",
       "        [44, 45, 46, 47]]])\n",
       "Coordinates:\n",
       "  * t        (t) datetime64[ns] 2018-08-15 2018-08-16 2018-08-17 2018-08-18\n",
       "  * x        (x) int64 0 1 2\n",
       "  * y        (y) int64 0 1 2 3"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da = xr.DataArray(np.arange(4*12).reshape((4,3, 4)), dims=['t', 'x', 'y'], \n",
    "                  coords={'t': time, 'x': [0, 1, 2], 'y': [0,1,2,3]})\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (t: 4)>\n",
       "array([ 0, 17, 47, 46])\n",
       "Coordinates:\n",
       "  * t        (t) datetime64[ns] 2018-08-15 2018-08-16 2018-08-18 2018-08-18\n",
       "    x        (t) int64 0 1 2 2\n",
       "    y        (t) int64 0 1 3 2"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "da.sel(t=t, x=ix, y=iy, method='nearest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "\"not all values found in index 't'\"",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-36-ebfdffc0b00e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mda\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mix\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0miy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'nearest'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtolerance\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimedelta\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdays\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/builds/anaconda/envs/xarray_stable/lib/python3.6/site-packages/xarray/core/dataarray.py\u001b[0m in \u001b[0;36msel\u001b[0;34m(self, indexers, method, tolerance, drop, **indexers_kwargs)\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0mindexers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meither_dict_or_kwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexers_kwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'sel'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    783\u001b[0m         ds = self._to_temp_dataset().sel(\n\u001b[0;32m--> 784\u001b[0;31m             indexers=indexers, drop=drop, method=method, tolerance=tolerance)\n\u001b[0m\u001b[1;32m    785\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_from_temp_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    786\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/builds/anaconda/envs/xarray_stable/lib/python3.6/site-packages/xarray/core/dataset.py\u001b[0m in \u001b[0;36msel\u001b[0;34m(self, indexers, method, tolerance, drop, **indexers_kwargs)\u001b[0m\n\u001b[1;32m   1507\u001b[0m         \u001b[0mindexers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meither_dict_or_kwargs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexers_kwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'sel'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1508\u001b[0m         pos_indexers, new_indexes = remap_label_indexers(\n\u001b[0;32m-> 1509\u001b[0;31m             self, indexers=indexers, method=method, tolerance=tolerance)\n\u001b[0m\u001b[1;32m   1510\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpos_indexers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdrop\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdrop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1511\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_replace_indexes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnew_indexes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/builds/anaconda/envs/xarray_stable/lib/python3.6/site-packages/xarray/core/coordinates.py\u001b[0m in \u001b[0;36mremap_label_indexers\u001b[0;34m(obj, indexers, method, tolerance, **indexers_kwargs)\u001b[0m\n\u001b[1;32m    353\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    354\u001b[0m     pos_indexers, new_indexes = indexing.remap_label_indexers(\n\u001b[0;32m--> 355\u001b[0;31m         \u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv_indexers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmethod\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtolerance\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtolerance\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    356\u001b[0m     )\n\u001b[1;32m    357\u001b[0m     \u001b[0;31m# attach indexer's coordinate to pos_indexers\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/builds/anaconda/envs/xarray_stable/lib/python3.6/site-packages/xarray/core/indexing.py\u001b[0m in \u001b[0;36mremap_label_indexers\u001b[0;34m(data_obj, indexers, method, tolerance)\u001b[0m\n\u001b[1;32m    248\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    249\u001b[0m             idxr, new_idx = convert_label_indexer(index, label,\n\u001b[0;32m--> 250\u001b[0;31m                                                   dim, method, tolerance)\n\u001b[0m\u001b[1;32m    251\u001b[0m             \u001b[0mpos_indexers\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0midxr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    252\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mnew_idx\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/builds/anaconda/envs/xarray_stable/lib/python3.6/site-packages/xarray/core/indexing.py\u001b[0m in \u001b[0;36mconvert_label_indexer\u001b[0;34m(index, label, index_name, method, tolerance)\u001b[0m\n\u001b[1;32m    187\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexer\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    188\u001b[0m                 raise KeyError('not all values found in index %r'\n\u001b[0;32m--> 189\u001b[0;31m                                % index_name)\n\u001b[0m\u001b[1;32m    190\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnew_index\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: \"not all values found in index 't'\""
     ]
    }
   ],
   "source": [
    "da.sel(t=t, x=ix, y=iy, method='nearest', tolerance=dt.timedelta(days=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "datetime.timedelta(0, 43200)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dt.timedelta(days=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
